import time
from fastapi import Request, Response
from utils.logger import logger
from utils.logger import TraceID

# 定义中间件功能
async def add_process_time_header(request: Request, call_next):
    start_time = time.time()
    response = await call_next(request)
    process_time = round((time.time() - start_time) * 1000, 2)
    response.headers["X-Process-Time"] = f'{str(process_time)}ms'
    logger.info(f"{request.method} {request.url} {response.status_code} {process_time}ms")
    return response

 # 设置日志的全链路追踪
async def set_logger(request: Request, call_next):
    REQUEST_ID_KEY = "X-Request-Id"
    _req_id_val = request.headers.get(REQUEST_ID_KEY, "1")
    req_id = TraceID.set(_req_id_val)
    response: Response = await call_next(request)
    response.headers[REQUEST_ID_KEY] = req_id.get()
    return response